library(tidyverse)
library(rstan)
options(mc.cores = parallel::detectCores())
rstan_options(auto_write = T)

args <- commandArgs(trailingOnly = TRUE)

if (!any(args == "stan_data")) {
  stop("Provide a data input with the keyword stan_data")
} else if (!any(args == "saveto")) {
  stop("Provide an output filename with the keyword saveto")
} else {
  stan_data_idx <- which(args == "stan_data")
  stan_data <- readRDS(args[stan_data_idx + 1])
  
  saveto_idx <- which(args == "saveto")
  filename <- args[saveto_idx + 1]
  
  if (!any(args == "seed")) {
    seed <- 1
  } else {
    seed_idx <- which(args == "seed")
    seed <- as.integer(args[seed_idx + 1])
  }
  
  print(Sys.time())
  
  system.time(m <- stan_model(model_code = read_file("Scripts/Poisson_model_age_sex_effect_nu_shrunk.stan")))
  
  vb_fit <- vb(object = m, data = stan_data, iter = 50000, seed = seed)
  
  print(paste("saving to", filename))
  
  res <- list(vb_fit = vb_fit, data = stan_data)
  
  saveRDS(res, file = paste(filename, "rds", sep = "."))
  
  print(Sys.time())
}
